#include "T2KParticle.h"

class NeutrinoUtils
{
 public:
  enum NeutrinoParent{
    kNull = -999,
    kBad = -2,
    kLost = -1,
    kOther=0,
    kPion,
    kKaon,
    kMuon
  };
  
  static void IniCC();
  static void IniCCQE();

  //--- Reconstruction related
  static const TLorentzVector * fNeutrinoRec;
  static T2KParticle * fMuonRec; 
  static T2KParticle * fProtonRec; 
  static T2KParticle * fXRec;
  static Double_t fDeltaPTTSim;
  static Double_t fDeltaPTTRec;
  static Double_t fSumPtSim;
  static Double_t fSumPtRec;
  static Double_t fMuonToyKt;
  static Double_t fDelta2PlusMassSim;
  static Double_t fDelta2PlusMassRec;

  static const TLorentzVector * fNeutrinoGuess;
  static T2KParticle * fMuonGuess;

  //--- Simulation related
  static const TLorentzVector * fNeutrinoSim;
  static T2KParticle * fMuonSim;
  static T2KParticle * fProtonSim;
  static T2KParticle * fXSim;
  static T2KParticle * fNeutronToy;
  static T2KParticle * fProtonToy;
  static T2KParticle * fProtonDet;

  static const TVector3 * fNeutrinoParentDecPointSim;
  
  static Double_t fOsciMuMuProb;
  static Int_t fkOsciMuSurvive;
  static Int_t fNeutrinoParentPDG;
  static Int_t fNeutrinoParentType;
  static Int_t fNeutrinoType;
  static Int_t fNeutMode;
  static Int_t fTopology;
  static Int_t fSelX_Charge;
  static Int_t fSelX_PID;
  static Int_t fEventTypeSim;
  static Int_t fSimOtherPDG;
  static Int_t fNECalTracks;
  static Int_t fECalPDGSim;

  static void SetMuonNRec(const Int_t var){ fMuonNRec = var; }
  static Int_t GetMuonNRec(){return fMuonNRec;}
  static void SetMuonNSim(const Int_t var){ fMuonNSim = var; }
  static Int_t GetMuonNSim(){return fMuonNSim;}
  
  static void SetProtonNRec(const Int_t var){ fProtonNRec = var; }
  static Int_t GetProtonNRec(){return fProtonNRec;}
  static void SetProtonNSim(const Int_t var){ fProtonNSim = var; }
  static Int_t GetProtonNSim(){return fProtonNSim;}

  static void SetXNRec(const Int_t var){ fXNRec = var; }
  static Int_t GetXNRec(){return fXNRec;}
  static void SetXNSim(const Int_t var){ fXNSim = var; }
  static Int_t GetXNSim(){return fXNSim;}

  static void SetMultiplicity(const Int_t mul){fMultiplicity = mul;}
  static Int_t GetMultiplicity(){return fMultiplicity;}
  
  //only rec. meaningful; phi angle in the transver plane w.r.t. beam direction, -180 ~ 180, by definition muon at (-)180.
  static Double_t fBetaRec[10];  
  static Double_t fGammaRec[10];
  
  //-------- Calculations ---------
  //--- Reconstruction related
  Double_t CalcBetaRec(const Double_t dir0, const Double_t dir1, const Double_t dir2);
  Double_t CalcGammaRec(const Double_t dir0, const Double_t dir1, const Double_t dir2);
 
  static void SetRecGuessKinematics(const Bool_t kdummy=kFALSE);

  //---Simulation related
  static void MuonNSimSanityCheck();

  static void SetSimKinematics();
  static void SetTransverseKinematics();
  static void SetDelta2PlusMass();

  //-------- IO ---------
  static TTree * GetTree();
  static TList * GetHistList();

  static Int_t FillCount(TList *lin, const Int_t var){ return NeutrinoTools::FillTH1I(lin, kcount, var);}

 private:
  //==========================================================
  //Variables and setters
  //==========================================================
  //--- Reconstruction related
  static Double_t fdNeutrinoGuessNeutrinoSim;
  static Double_t fdNeutrinoGuessMuonRec;
  static const Int_t fCutNR;
  static Double_t fGuessCorrection;
  static void SetdNeutrinoGuessNeutrinoSim(); 
  static void SetdNeutrinoGuessMuonRec();

  static Int_t fMuonNRec;
  static Int_t fProtonNRec;
  static Int_t fXNRec;
  
  static Int_t fMultiplicity; 
  static Double_t fMeanBetaRec; 
  static void SetMeanBetaRec(const Bool_t kdummy = kFALSE);
  static Double_t fMeanGammaRec;//use Mean instead of sum to get rid of multiplicity
  static void SetMeanGammaRec(const Bool_t kdummy = kFALSE);

  ///---Simulation related
  static Bool_t IsCC(){ return TMath::Abs(fNeutMode)<30;}

  static Int_t fMuonNSim;
  static Int_t fProtonNSim;
  static Int_t fXNSim;
 
  enum histID{
    kcount,
    kDUMMY
  };

};

//__________________________________________________________

const TLorentzVector * NeutrinoUtils::fNeutrinoRec = 0x0;
T2KParticle *NeutrinoUtils::fMuonRec=0x0; 
T2KParticle *NeutrinoUtils::fProtonRec=0x0; 
T2KParticle *NeutrinoUtils::fXRec=0x0; 
Double_t NeutrinoUtils::fDeltaPTTSim = NeutrinoTools::kIniValue;
Double_t NeutrinoUtils::fDeltaPTTRec = NeutrinoTools::kIniValue;
Double_t NeutrinoUtils::fSumPtSim = NeutrinoTools::kIniValue;
Double_t NeutrinoUtils::fSumPtRec = NeutrinoTools::kIniValue;
Double_t NeutrinoUtils::fMuonToyKt = NeutrinoTools::kIniValue;
Double_t NeutrinoUtils::fDelta2PlusMassSim = NeutrinoTools::kIniValue;
Double_t NeutrinoUtils::fDelta2PlusMassRec = NeutrinoTools::kIniValue;
const TLorentzVector * NeutrinoUtils::fNeutrinoGuess = 0x0;
T2KParticle *NeutrinoUtils::fMuonGuess=0x0; 
const TLorentzVector * NeutrinoUtils::fNeutrinoSim = 0x0;
T2KParticle *NeutrinoUtils::fMuonSim=0x0; 
T2KParticle *NeutrinoUtils::fProtonSim=0x0; 
T2KParticle *NeutrinoUtils::fXSim=0x0; 
T2KParticle *NeutrinoUtils::fNeutronToy=0x0; 
T2KParticle *NeutrinoUtils::fProtonToy=0x0; 
T2KParticle *NeutrinoUtils::fProtonDet=0x0; 

const TVector3 * NeutrinoUtils::fNeutrinoParentDecPointSim = 0x0;

Double_t NeutrinoUtils::fOsciMuMuProb = NeutrinoTools::kIniValue;
Int_t NeutrinoUtils::fkOsciMuSurvive = NeutrinoTools::kIniValue;
Int_t NeutrinoUtils::fNeutrinoParentPDG = NeutrinoUtils::kNull;
Int_t NeutrinoUtils::fNeutrinoParentType = NeutrinoUtils::kNull;
Int_t NeutrinoUtils::fNeutrinoType = NeutrinoTools::kIniValue;
Int_t NeutrinoUtils::fNeutMode = NeutrinoTools::kIniValue;
Int_t NeutrinoUtils::fTopology = NeutrinoTools::kIniValue;
Int_t NeutrinoUtils::fSelX_Charge = NeutrinoTools::kIniValue;
Int_t NeutrinoUtils::fSelX_PID = NeutrinoTools::kIniValue;
Int_t NeutrinoUtils::fEventTypeSim = NeutrinoTools::kIniValue;
Int_t NeutrinoUtils::fSimOtherPDG = NeutrinoTools::kIniValue;
Int_t NeutrinoUtils::fNECalTracks = NeutrinoTools::kIniValue;
Int_t NeutrinoUtils::fECalPDGSim = NeutrinoTools::kIniValue;

Int_t NeutrinoUtils::fMuonNRec = NeutrinoTools::kIniValue;
Int_t NeutrinoUtils::fMuonNSim = NeutrinoTools::kIniValue;
Int_t NeutrinoUtils::fProtonNRec = NeutrinoTools::kIniValue;
Int_t NeutrinoUtils::fProtonNSim = NeutrinoTools::kIniValue;
Int_t NeutrinoUtils::fXNRec = NeutrinoTools::kIniValue;
Int_t NeutrinoUtils::fXNSim = NeutrinoTools::kIniValue;

Int_t NeutrinoUtils::fMultiplicity = NeutrinoTools::kIniValue;

Double_t NeutrinoUtils::fBetaRec[10]={NeutrinoTools::kIniValue, NeutrinoTools::kIniValue, NeutrinoTools::kIniValue, NeutrinoTools::kIniValue, NeutrinoTools::kIniValue, NeutrinoTools::kIniValue,NeutrinoTools::kIniValue, NeutrinoTools::kIniValue, NeutrinoTools::kIniValue,NeutrinoTools::kIniValue};
Double_t NeutrinoUtils::fMeanBetaRec = NeutrinoTools::kIniValue;
Double_t NeutrinoUtils::fGammaRec[10]={NeutrinoTools::kIniValue, NeutrinoTools::kIniValue, NeutrinoTools::kIniValue, NeutrinoTools::kIniValue, NeutrinoTools::kIniValue, NeutrinoTools::kIniValue,NeutrinoTools::kIniValue, NeutrinoTools::kIniValue, NeutrinoTools::kIniValue,NeutrinoTools::kIniValue};
Double_t NeutrinoUtils::fMeanGammaRec = NeutrinoTools::kIniValue;

Double_t NeutrinoUtils::fdNeutrinoGuessNeutrinoSim = NeutrinoTools::kIniValue;
Double_t NeutrinoUtils::fdNeutrinoGuessMuonRec = NeutrinoTools::kIniValue;

Double_t NeutrinoUtils::fGuessCorrection = 1;
const Int_t NeutrinoUtils::fCutNR = 1;

//_______________________________________________________________________________________________
//_______________________________________________________________________________________________
void NeutrinoUtils::IniCC()
{
  delete fMuonSim; fMuonSim= new T2KParticle;
  delete fMuonRec; fMuonRec= new T2KParticle;
  delete fMuonGuess; fMuonGuess= new T2KParticle;

  fMuonNRec = 0;
  fMuonNSim = 0;
}

void NeutrinoUtils::IniCCQE()
{
  IniCC();
  delete fProtonSim; fProtonSim= new T2KParticle;
  delete fXSim;    fXSim= new T2KParticle;
  delete fNeutronToy; fNeutronToy= new T2KParticle;
  delete fProtonToy; fProtonToy= new T2KParticle;
  delete fProtonDet; fProtonDet= new T2KParticle;
  delete fProtonRec; fProtonRec= new T2KParticle;
  delete fXRec;   fXRec= new T2KParticle;
  fDeltaPTTSim = NeutrinoTools::kIniValue;
  fDeltaPTTRec = NeutrinoTools::kIniValue;
  fSumPtSim = NeutrinoTools::kIniValue;
  fSumPtRec = NeutrinoTools::kIniValue;
  fMuonToyKt = NeutrinoTools::kIniValue;
  fDelta2PlusMassSim = NeutrinoTools::kIniValue;
  fDelta2PlusMassRec = NeutrinoTools::kIniValue;
  fProtonNRec = 0;
  fProtonNSim = 0;
  fXNRec = 0;
  fXNSim = 0;
}

void NeutrinoUtils::SetRecGuessKinematics(const Bool_t kdummy)
{ 
  SetdNeutrinoGuessMuonRec();
  SetMeanBetaRec(kdummy);
  SetMeanGammaRec(kdummy);
}

void NeutrinoUtils::SetSimKinematics()
{ 
  SetdNeutrinoGuessNeutrinoSim();
}

void NeutrinoUtils::SetTransverseKinematics()
{
  fMuonSim->SetTransverseKinematics(fProtonSim);
  fProtonSim->SetTransverseKinematics(fMuonSim);
  fXSim->SetTransverseKinematics(fMuonSim);

  fMuonRec->SetTransverseKinematics(fProtonRec);
  fProtonRec->SetTransverseKinematics(fMuonRec);
  fXRec->SetTransverseKinematics(fMuonRec);

  fProtonToy->SetTransverseKinematics(fMuonSim);
  fProtonDet->SetTransverseKinematics(fMuonSim);

  //------------
  fDeltaPTTSim = NeutrinoTools::kIniValue;
  if(fProtonSim->fPt && fXSim->fPt){
    fDeltaPTTSim = NeutrinoTools::GetDeltaPTT(fProtonSim->fKt, fXSim->fKt);
  }
  fDeltaPTTRec = NeutrinoTools::kIniValue;
  if(fProtonRec->fPt && fXRec->fPt){
    fDeltaPTTRec = NeutrinoTools::GetDeltaPTT(fProtonRec->fKt, fXRec->fKt);
  }
  //------------
  fSumPtSim = NeutrinoTools::kIniValue;
  if(fProtonSim->fPt && fMuonSim->fPt){
    fSumPtSim = NeutrinoTools::GetSumPt(fProtonSim->fPt, fMuonSim->fPt);
  }
  fSumPtRec = NeutrinoTools::kIniValue;
  if(fProtonRec->fPt && fMuonRec->fPt){
    fSumPtRec = NeutrinoTools::GetSumPt(fProtonRec->fPt, fMuonRec->fPt);
  }
  //=========

  fMuonToyKt = NeutrinoTools::kIniValue;
  if(fMuonSim->fLorentzVec && fProtonToy->fLorentzVec){
    fMuonToyKt = NeutrinoTools::GetPTT(fMuonSim->fLorentzVec->Vect(), fNeutrinoSim->Vect(), fProtonToy->fLorentzVec->Vect());
  }
}

void NeutrinoUtils::SetDelta2PlusMass()
{
  fDelta2PlusMassSim = NeutrinoTools::GetMass(fProtonSim->fLorentzVec, fXSim->fLorentzVec);
  fDelta2PlusMassRec = NeutrinoTools::GetMass(fProtonRec->fLorentzVec, fXRec->fLorentzVec);
}

void NeutrinoUtils::SetdNeutrinoGuessNeutrinoSim()
{
  //flat in cos, not tha angle
  //tested with = fNeutrinoRec; should be 0; confirmed
  if(!fNeutrinoGuess){
    fdNeutrinoGuessNeutrinoSim = NeutrinoTools::kNullPointer;
  }
  else{
    const Double_t tmpcos= NeutrinoTools::GetCos(fNeutrinoGuess, fNeutrinoSim, "NeutrinoUtils::SetdNeutrinoGuessNeutrinoSimGetCos");
    if(tmpcos== NeutrinoTools::kZeroDivider){
      fdNeutrinoGuessNeutrinoSim = NeutrinoTools::kZeroDivider;
    }
    else{
      fdNeutrinoGuessNeutrinoSim =  tmpcos - 1;
    }
  }
}

void NeutrinoUtils::SetdNeutrinoGuessMuonRec()
{
  //flat in cos, not tha angle

  //calculation confirmed by
  //tree->Draw("fdNeutrinoGuessMuonRec","fMultiplicity==1") //should = 1
  
  //tested with = &(fMuonRec->Vect())// should be 0; confirmed
  if(!fMuonRec->fLorentzVec || !fNeutrinoGuess){
    fdNeutrinoGuessMuonRec = NeutrinoTools::kNullPointer;
  }
  else{
    const Double_t tmpcos = NeutrinoTools::GetCos(fNeutrinoGuess, fMuonRec->fLorentzVec, "NeutrinoUtils::SetdNeutrinoGuessMuonRecGetCos");
    if(tmpcos== NeutrinoTools::kZeroDivider){
      fdNeutrinoGuessMuonRec = NeutrinoTools::kZeroDivider;
    }
    else{
      fdNeutrinoGuessMuonRec = tmpcos - 1;
    }
  }
}

void NeutrinoUtils::SetMeanBetaRec(const Bool_t kdummy)
{
  if(kdummy){
    fMeanBetaRec = NeutrinoTools::kIniValue;
    return;
  }

  fMeanBetaRec = 0;
  for(Int_t ipar=0; ipar<fMultiplicity; ipar++){
    //remove muon contribution
    if(ipar == fMuonRec->fID){
      continue;
    }

    fMeanBetaRec += fBetaRec[ipar];
  }

  if(fMultiplicity==1){
    fMeanBetaRec=NeutrinoTools::kZeroDivider;
  }
  else{//multiplicity >= 2
    fMeanBetaRec /= (fMultiplicity-1);
  }
}

void NeutrinoUtils::SetMeanGammaRec(const Bool_t kdummy)
{
  if(kdummy){
    fMeanGammaRec = NeutrinoTools::kIniValue;
    return;
  }

  fMeanGammaRec = 0;
  for(Int_t ipar=0; ipar<fMultiplicity; ipar++){
    //remove muon contribution
    if(ipar == fMuonRec->fID){
      continue;
    }
    fMeanGammaRec += fGammaRec[ipar];
  }

  if(fMultiplicity==1){
    fMeanGammaRec = NeutrinoTools::kZeroDivider;
  }
  else{
    fMeanGammaRec /= (fMultiplicity-1);
  }
}

void NeutrinoUtils::MuonNSimSanityCheck()
{
  //sanity check -1
  if(fMuonNSim==0 && IsCC()){
    printf("MuonNSimSanityCheck wrong fMuonNSim CC %d %d\n", fMuonNSim, fNeutMode); exit(1);
  }
  //sanity check -2
  if(fMuonNSim>0 && !IsCC()){
    printf("MuonNSimSanityCheck wrong fMuonNSim NC %d %d\n", fMuonNSim, fNeutMode); exit(1);
  }
}

Double_t NeutrinoUtils::CalcBetaRec(const Double_t dir0, const Double_t dir1, const Double_t dir2)
{
  //
  //magnitude does not magger
  //w.r.t. beam direction (z-axis) and muon opposite direction  (x-axis)
  //-180 ~ 180
  //
  if(!fMuonRec->fPt){
    return NeutrinoTools::kNullPointer;
  }

  const TVector3 *vt = NeutrinoTools::GetVecT(fNeutrinoRec, dir0, dir1, dir2);

  //muon pt direction := -180 deg, i.e. -x axis
  const TVector3 xaxis = (*fMuonRec->fPt)*(-1);

  //0 ~ 180

  
  Double_t beta = NeutrinoTools::GetAngle(&xaxis, vt, "NeutrinoUtils::CalcBetaRec");

  if( (fNeutrinoRec->Vect()).Dot(xaxis.Cross(*vt)) < 0 )
    beta *= -1;

  if(beta<-180 + 1e-3){
    beta+= 360;
  }

  delete vt;
  return beta;
}

Double_t NeutrinoUtils::CalcGammaRec(const Double_t dir0, const Double_t dir1, const Double_t dir2)
{
  const TVector3 vec(dir0, dir1, dir2);

  const TVector3 refdir=fNeutrinoRec->Vect();
  return NeutrinoTools::GetAngle(&vec, &refdir, "NeutrinoUtils::CalcGammaRec");
}

//___________________________________________________________________________________

//___________________________________________________________________________________________________

TTree * NeutrinoUtils::GetTree()
{
  TTree * tout = new TTree("tree","tree");

  const Int_t spl = 1;

  fMuonRec->SetTree("fMuonRec",tout,spl);
  fMuonSim->SetTree("fMuonSim",tout,spl);

  if(fProtonRec && fProtonSim){
    fProtonRec->SetTree("fProtonRec",tout,spl);
    fProtonSim->SetTree("fProtonSim",tout,spl);
  }

  if(fProtonToy){
    fNeutronToy->SetTree("fNeutronToy",tout,spl);
    fProtonToy->SetTree("fProtonToy",tout,spl);
  }

  if(fProtonDet){
    fProtonDet->SetTree("fProtonDet",tout,spl);
  }

  if(fXRec && fXSim){
    fXRec->SetTree("fXRec",tout,spl);
    fXSim->SetTree("fXSim",tout,spl);
  }

  tout->Branch("fDeltaPTTSim",&fDeltaPTTSim,"fDeltaPTTSim/D");
  tout->Branch("fDeltaPTTRec",&fDeltaPTTRec,"fDeltaPTTRec/D");
  tout->Branch("fSumPtSim",&fSumPtSim,"fSumPtSim/D");
  tout->Branch("fSumPtRec",&fSumPtRec,"fSumPtRec/D");
  tout->Branch("fMuonToyKt",&fMuonToyKt,"fMuonToyKt/D");
  tout->Branch("fDelta2PlusMassSim",&fDelta2PlusMassSim,"fDelta2PlusMassSim/D");
  tout->Branch("fDelta2PlusMassRec",&fDelta2PlusMassRec,"fDelta2PlusMassRec/D");
  
  tout->Branch("fNeutrinoSim","TLorentzVector", &fNeutrinoSim, 128000, spl);
  tout->Branch("fNeutrinoRec", "TVector3", &(fMuonRec->fNuDir), 128000, spl);
  tout->Branch("fNeutrinoGuess", "TVector3", &(fMuonGuess->fNuDir), 128000, spl);
  
  tout->Branch("fdNeutrinoGuessNeutrinoSim", &fdNeutrinoGuessNeutrinoSim, "fdNeutrinoGuessNeutrinoSim/D");
  tout->Branch("fdNeutrinoGuessMuonRec", &fdNeutrinoGuessMuonRec, "fdNeutrinoGuessMuonRec/D");

  tout->Branch("fGuessCorrection", &fGuessCorrection, "fGuessCorrection/D");
  tout->Branch("fMuonPtGuess", &(fMuonGuess->fPtMag), "fMuonPtGuess/D");

  tout->Branch("fOsciMuMuProb",&fOsciMuMuProb,"fOsciMuMuProb/D");
  tout->Branch("fkOsciMuSurvive",&fkOsciMuSurvive,"fkOsciMuSurvive/I");
  tout->Branch("fNeutrinoParentPDG",&fNeutrinoParentPDG,"fNeutrinoParentPDG/I");
  tout->Branch("fNeutrinoParentType",&fNeutrinoParentType,"fNeutrinoParentType/I");
  tout->Branch("fNeutrinoType",&fNeutrinoType,"fNeutrinoType/I");
  tout->Branch("fNeutMode",&fNeutMode,"fNeutMode/I");
  tout->Branch("fTopology",&fTopology,"fTopology/I");
  tout->Branch("fSelX_Charge",&fSelX_Charge,"fSelX_Charge/I");
  tout->Branch("fSelX_PID",&fSelX_PID,"fSelX_PID/I");
  tout->Branch("fEventTypeSim",&fEventTypeSim,"fEventTypeSim/I");
  tout->Branch("fSimOtherPDG",&fSimOtherPDG,"fSimOtherPDG/I");
  tout->Branch("fNECalTracks",&fNECalTracks,"fNECalTracks/I");
  tout->Branch("fECalPDGSim",&fECalPDGSim,"fECalPDGSim/I"); 
  tout->Branch("fNeutrinoParentDecPointSim", "TVector3", &fNeutrinoParentDecPointSim, 128000, spl);

  tout->Branch("fMultiplicity",&fMultiplicity,"fMultiplicity/I");

  tout->Branch("fMuonNSim",&fMuonNSim,"fMuonNSim/I");
  tout->Branch("fMuonNRec",&fMuonNRec,"fMuonNRec/I");
  tout->Branch("fProtonNSim",&fProtonNSim,"fProtonNSim/I");
  tout->Branch("fProtonNRec",&fProtonNRec,"fProtonNRec/I");
  tout->Branch("fXNSim",&fXNSim,"fXNSim/I");
  tout->Branch("fXNRec",&fXNRec,"fXNRec/I");

  tout->Branch("fMeanBetaRec", &fMeanBetaRec, "fMeanBetaRec/D");
  tout->Branch("fMeanGammaRec", &fMeanGammaRec, "fMeanGammaRec/D");

  return tout;
}

TList * NeutrinoUtils::GetHistList()
{
  TList *lout=new TList;
  lout->SetName("lout");
  
  TH1D * hcount = new TH1D("count","",30, -0.5, 29.5); lout->AddAt(hcount,kcount);  
  return lout;
}
